import numpy as np
import os
import cv2
import matplotlib.pyplot as plt
import json
import tempfile
import img_transform

data_dir = "/Users/pitaloveu/Downloads/pic"
anno_file = "/Users/pitaloveu/Downloads/coor.txt"

def random_synthesise( wrapper_json_dir , wrapper_img_dir , content_dir , N , \
        img_storage ):
    """
    params:
        wrapper_json_dir contains all the images annotations for wrapping
        wrapper_img_dir contains all the images for wrapping
        content_dir contains all the images as content

    algorithms:
        the core function will be loaded actually N times

    return a list containing N numbers of synthesised images
    """
    anno_json_list = os.listdir( wrapper_json_dir )
    anno_json_list = [ x for x in anno_json_list if x[-5:] == ".json" ]

    anno_dict_list = []

    for anno_json in anno_json_list:
        with open( os.path.join( wrapper_json_dir , anno_json )  , "r" ) as json_data:
            d = json.load( json_data )

            if d['labeled'] == False or len( d['outputs']['object'] ) == 0 :
                continue

            # some time the json data will contain some bezier curves
            # the bezier ones will be ignored because we can't handle this type
            has_cubic_bezier = False

            for obj in d['outputs']['object']:
                if "cubic_bezier" in obj:
                    has_cubic_bezier = True
                    break
            if has_cubic_bezier:
                continue
            
            json_dict = {}

            img_name = d['path'].strip().split("\\")[-1]
            json_dict["img_path"] = os.path.join( wrapper_img_dir , img_name )

            for obj in d['outputs']['object']:
                if "bndbox" in obj:
                    xmin = obj['bndbox']['xmin']
                    ymin = obj['bndbox']['ymin']
                    xmax = obj['bndbox']['xmax']
                    ymax = obj['bndbox']['ymax']

                    json_dict[ obj['name'] ] = [ (xmin, ymin) , (xmax, ymin) , (xmax, ymax) , (xmin, ymax) ]

                if "polygon" in obj:
                    x1 = obj['polygon']['x1']
                    y1 = obj['polygon']['y1']
                    x2 = obj['polygon']['x2']
                    y2 = obj['polygon']['y2']
                    x3 = obj['polygon']['x3']
                    y3 = obj['polygon']['y3']
                    x4 = obj['polygon']['x4']
                    y4 = obj['polygon']['y4']

                    LU, RU, RD, LD = (), (), (), ()

                    pt1 = ( x1, y1 )
                    pt2 = ( x2, y2 )
                    pt3 = ( x3, y3 )
                    pt4 = ( x4, y4 )

                    pts = [ pt1 , pt2 , pt3 , pt4 ]

                    centra_x = ( x1 + x2 + x3 + x4 ) / 4
                    centra_y = ( y1 + y2 + y3 + y4 ) / 4

                    for pt in [ pt1 , pt2 , pt3 , pt4 ]:
                        if pt[0] < centra_x and pt[1] < centra_y :
                            LU = pt
                        if pt[0] > centra_x and pt[1] < centra_y :
                            RU = pt
                        if pt[0] > centra_x and pt[1] > centra_y :
                            RD = pt
                        if pt[0] < centra_x and pt[1] > centra_y :
                            LD = pt

                    if LU == [] or RU == [] or RD == [] or LD == []:
                        continue
                    json_dict[ obj['name'] ] = [ LU , RU , RD , LD ]

        anno_dict_list.append( json_dict )

    # reading the content into a img list
    content_img_list = os.listdir( content_dir )
    content_img_list = [ os.path.join( content_dir , x ) for x in content_img_list ]

    # generate N random numbers 
    wrapper_indexes = np.random.randint( low = 0 , high = len( anno_dict_list ) , size = N )
    #wrapper_indexes = list( range( len(anno_dict_list) ) )
    #print( wrapper_indexes )
    content_indexes = np.random.randint( low = 0 , high = len( content_img_list ) , size = N )
    #content_indexes = np.random.randint( low = 0 , high = len( anno_dict_list ) , size = len( anno_dict_list ) )

    for index1, index2 in zip( wrapper_indexes , content_indexes ):
        d = anno_dict_list[index1]
        anno_img_path = d['img_path']

        #print( index1 )
        if not os.path.exists( anno_img_path ):
            continue
        #print( anno_img_path )
        content_img_path = content_img_list[index2]

        wrapping_img = cv2.imread( anno_img_path )
        content_img  = cv2.imread( content_img_path )


        wrapping_img , dl , ds = img_transform.get_random_wrapper( wrapping_img , d['l'] , d['s'])
        syn_img = img_transform.img_synthesis( wrapping_img , content_img , ds )
        #syn_img = img_synthesis( wrapping_img , content_img , d['s'] )
        img_name = tempfile.mktemp( suffix = ".jpg" , dir = img_storage )

        #plt.imshow( syn_img )
        #plt.show()
        cv2.imwrite( img_name , syn_img )

if __name__ == "__main__":

    anno_dir = "/Users/pitaloveu/Desktop/images/outputs"
    img_base_dir = "/Users/pitaloveu/Desktop/images"
    content_dir = "/Users/pitaloveu/Desktop/image_content_for_synthesise"

    img_storage = "/Users/pitaloveu/Desktop/generated_imgs"

    random_synthesise( anno_dir , img_base_dir , content_dir , 600 , img_storage )
